#import numpy as np
from Env.action import *
import random


tpdic = {'坦克':0 , '无人机':1 , '士兵': 2}
class Generator:
    def __init__(self):
        self.last_reward = 0

    def generate(self, raw_obs):
        ids = self.gen_ids(raw_obs)
        #print(ids)
        valid_action = {}
        meta_action = random.choice([0, 1])
        command = []
        #print('---------------\n',action['target'])
        if meta_action == 0:
            if len(ids['tgt_id']):
                tid = random.choice(ids['tgt_id'])
                for uid in ids['uac_id']:
                    if random.random()>0.75:
                        continue
                    cmd = FireTargetWrc(uid, tid)
                    command.append(cmd)
        else:
            tx = random.randint(1,100)
            ty = random.randint(1,100)
            x = tx * 30 + 1000
            y = tx * 30 + 6000
            for uid in ids['uac_id']:
                if random.random()>0.8:
                    continue
                cmd = MoveToPointWrc(uid, x, y, 50)
                command.append(cmd)
        return command

    def collect_features(self, raw_obs, env_step_info):
        feature_dic = self.gen_feature(raw_obs)
        feature_template_values = {
            "tank_state": [
                {'info': tfeat} for tfeat in feature_dic['tank']
            ],
            "wrc_state": [
                {'info': tfeat} for tfeat in feature_dic['wrc']
            ],
            "wrj_state": [
                {'info': tfeat} for tfeat in feature_dic['wrj']
            ],
            "bbb_state": [
                {'info': tfeat} for tfeat in feature_dic['bbb']
            ],
            "qb_state": [
                {'info': tfeat} for tfeat in feature_dic['qb']
            ],
        }
        return feature_template_values

    def calculate_reward(self, raw_obs, env_step_info):
        #reward = env_step_info.info['reward']
        red= raw_obs.red
        red_reward = self.calc_red(red)

        blue = raw_obs.blue
        blue_reward= self.calc_blue(blue)
        step_reward = red_reward + blue_reward - self.last_reward
        self.last_reward = red_reward + blue_reward
        return step_reward

    def calc_red(self, red):
        reward = 0
        for unit in red.tanks:
            reward += 10 * (unit.health/100)
        for unit in red.wrjs:
            reward += 1 * (unit.health/100)
        for unit in red.bbb:
            reward += 4 * (unit.health/100)
        for unit in red.wrc:
            reward += 3 * (unit.health/100)
        return -reward


    def calc_blue(self, red):
        reward = 0
        for unit in red.tanks:
            reward += 10 * (unit.health/100)
        for unit in red.wrjs:
            reward += 1 * (unit.health/100)
        for unit in red.sb:
            reward += 2 * (unit.health/100)
        return reward

    def gen_feature(self, obs):
        def type_enum(tp):
            return tpdic[tp]

        side_info = obs.red
        tanks = [ [u.x, u.y, u.health] for u in side_info.tanks if u.isalive]
        wrcs = [ [u.x, u.y, u.health] for u in side_info.wrc if u.isalive]
        wrjs = [ [u.x, u.y, u.health] for u in side_info.wrjs if u.isalive]
        bbbs = [ [u.x, u.y, u.health] for u in side_info.bbb if u.isalive]
        qbs = [ [u.x, u.y, type_enum(u.type), u.speed, u.health] for u in side_info.qbs if u.isalive]
        feature_dic = {
            'tank': tanks,
            'wrc': wrcs,
            'wrj': wrjs,
            'bbb': bbbs,
            'qb': qbs
        }
        return feature_dic

    def gen_ids(self, obs):
        side_info = obs.red
        uac_id = [ u.uid for u in side_info.wrc if u.isalive]
        tgt_id = [ u.uid for u in side_info.qbs if u.isalive]
        feature_dic = {
            'uac_id': uac_id,
            'tgt_id': tgt_id
        }
        return feature_dic

    def reset(self, raw_obs, env_step_info):
        self.last_reward = 0
        super().reset(raw_obs, env_step_info)

